"""
@ GOS (Guidance Point Search) path planning in 2D worskspace
@ Author: Gang Xu
@ Date: 2024.01.12
@ Function: A fast path generation method
"""
import time
import math
import numpy as np
import matplotlib.pyplot as plt
from shapely.geometry import LineString, Polygon

from rvo import env
from rvo.config import envs_type
from draw.plt2d import draw_polygon_2d
from rvo.poly_obstacle import PolyObstacle, Vertice
from rvo.utils import left_of, point_line_dist, normalize, pos_in_polygons
from rvo.utils import l2norm, unit_normal_vector, determin_between_line, path_length, smooth


class GOS(object):
    def __init__(self):
        self.s_start = None     # 必须为tuple型
        self.s_goal = None      # 必须为tuple型
        self.x_range = None
        self.y_range = None
        self.radius = -1
        self.turning_radius = -1
        self.epsilon = 1e-5
        self.inflation = 0.01
        self.path = []
        self.searched = dict()
        self.checked_obs = dict()
        self.global_guide_point = None
        self.length = 0.0
        self.guide_points = []
        self.poly_obs = None     # 分别用多边形和栅格点集表示障碍物
        self.inflated_obs = []              # 考虑机器人大小后的障碍物膨胀
        self.inflated_polygons = []
        self.visible_graph = []             # 障碍物膨胀后的扩展点
        self.feasible_verts_dict = dict()
        self.path_verts = []
        self.check_time = 0.0
        self.search_time = 0.0
        self.init_time = 0.0

    def init_rob_info(self, s_start, s_goal, rob_radius, inflation=0.01):
        self.s_start = s_start
        self.s_goal = s_goal
        self.radius = rob_radius
        self.turning_radius = rob_radius
        self.inflation = inflation
        self.path = []
        self.path_verts = [Vertice(self.s_start)]
        self.searched = dict()
        self.checked_obs = dict()
        self.global_guide_point = None
        self.length = 0.0
        self.guide_points = []
        self.check_time = 0.0
        self.search_time = 0.0
        self.init_time = 0.0

    def init_env(self, obstacles, x_range, y_range):
        t1 = time.time()
        self.x_range = x_range
        self.y_range = y_range
        self.poly_obs = obstacles

        agent_rad = self.radius + self.inflation
        combinedRadius = agent_rad + self.inflation
        for i in range(len(self.poly_obs)):
            self.scale_polygon_vertices(self.poly_obs[i], combinedRadius)
        for i in range(len(self.visible_graph)):
            self.feasible_verts_dict.update(self.visible_graph[i].obstacle_dict)
            vertices = self.visible_graph[i].vertices_
            for j in range(len(vertices)):
                if self.x_range[0] < vertices[j].point_[0] < self.x_range[1] \
                        and self.y_range[0] < vertices[j].point_[1] < self.y_range[1]:
                    continue
                else:
                    vertices[j].is_disrupted = True
                    self.visible_graph[i].unreachable_vertices.append(vertices[j])
        t2 = time.time()
        self.init_time += t2 - t1

    def global_search(self):
        pos = np.array(self.s_start)
        pos_next = np.array(self.s_goal)
        self.feasible_verts_dict[self.s_start] = (-1, -1)
        self.feasible_verts_dict[self.s_goal] = (-1, -1)
        if pos_in_polygons(pos, self.inflated_obs) or pos_in_polygons(pos_next, self.inflated_obs):
            return []
        interseted_obs = self.check_all_intersect_obstacles(pos, pos_next, self.inflated_obs)
        path = [self.s_start]
        while True:
            # self.plot_path(path + [pos_next, self.s_goal])
            guide_point = self.gos(pos, pos_next, interseted_obs)
            if guide_point is None:
                print("Solve Failed, there may not be a path to avoid collision.")
                return []
            is_interseted = self.check_intersect_obstacles(pos, guide_point, self.inflated_obs)
            if is_interseted is None:
                path.append(guide_point)
                ob_id, vert_id = self.feasible_verts_dict[guide_point]
                if ob_id != -1:
                    self.path_verts.append(self.visible_graph[ob_id].vertices_[vert_id])

                if path[-1] == self.s_goal:
                    self.path_verts.append(Vertice(self.s_goal))
                    # self.path = path
                    # self.length = path_length(path)
                    self.path, self.length = smooth(path, self.inflated_obs)
                    return self.path
                if self.searched and path[-1] == self.searched[(self.s_start, self.s_goal)][:2]:
                    del self.searched[(self.s_start, self.s_goal)]
                pos = np.array(path[-1])
                if self.searched:
                    pos_next = np.array(self.searched[(self.s_start, self.s_goal)][:2])
                else:
                    pos_next = np.array(self.s_goal)
            else:
                pos_next = np.array(guide_point)
            interseted_obs = self.check_all_intersect_obstacles(pos, pos_next, self.inflated_obs)

    def scale_polygon_vertices(self, polygon, combinedRadius):
        """求解多边形障碍物与机器人的近似闵可夫斯基和的顶点"""
        vertices_pos = []
        graph_pos = []
        for vertice in polygon.vertices_:
            p1 = normalize(vertice.previous_.point_ - vertice.point_)
            p2 = normalize(vertice.next_.point_ - vertice.point_)
            p1p2 = p2 - p1
            nLeft, nRight = unit_normal_vector(p1p2)
            n_scale = combinedRadius / math.sin(math.acos(p1.dot(p2)) / 2)
            vertices_pos.append(vertice.point_ + n_scale * nRight)  # 障碍物的顶点逆时针表示，因此边的右侧即为向障碍物外部扩展
            graph_pos.append(vertice.point_ + (n_scale + self.inflation) * nRight)
        self.inflated_polygons.append(Polygon(vertices_pos))
        self.inflated_obs.append(PolyObstacle(shape_dict={'shape': 'polygon'}, id=polygon.id,
                                              tid='poly' + str(polygon.id), vertices=vertices_pos))
        self.visible_graph.append(PolyObstacle(shape_dict={'shape': 'polygon'}, id=polygon.id,
                                               tid='poly' + str(polygon.id), vertices=graph_pos))

    def check_intersect_obstacles(self, pos, pos_next, sorted_obs, danger_obs=None):
        t1 = time.time()
        for obstacle in sorted_obs:
            if obstacle.is_danger_and_avoided or (danger_obs is not None and danger_obs.id == obstacle.id):
                continue
            line = LineString([pos, pos_next])
            if line.intersects(self.inflated_polygons[obstacle.id]):
                t2 = time.time()
                self.check_time += t2 - t1
                return obstacle
        t2 = time.time()
        self.check_time += t2 - t1
        return None

    def check_all_intersect_obstacles(self, pos, pos_next, sorted_obs):
        t1 = time.time()
        all_intersected_obs = []
        for obstacle in sorted_obs:
            if obstacle.is_danger_and_avoided:
                continue
            line = LineString([pos, pos_next])
            if line.intersects(self.inflated_polygons[obstacle.id]):
                all_intersected_obs.append(obstacle)
        t2 = time.time()
        self.check_time += t2 - t1
        return all_intersected_obs

    def determin_cand_verts(self, pos, pos_next, intersect_ob):
        idx = intersect_ob.id
        vertices = self.visible_graph[idx].vertices_
        max_left_vert, max_right_vert = None, None
        max_left_dist, max_right_dist = -1, -1
        disrupted_verts = {'left_side': [], 'right_side': []}
        for v in vertices:
            v_in_side = left_of(pos, pos_next, v.point_)
            r = determin_between_line(pos, pos_next, v.point_)
            if r and v.convex_ and not v.is_disrupted:
                dist = point_line_dist(pos, pos_next, v.point_)
                if v_in_side > 0:
                    if dist > max_left_dist:
                        max_left_dist = dist
                        max_left_vert = v
                elif v_in_side < 0:
                    if dist > max_right_dist:
                        max_right_dist = dist
                        max_right_vert = v
            elif v.is_disrupted:
                if v_in_side > 0:
                    disrupted_verts['left_side'].append(v.id_)
                elif v_in_side < 0:
                    disrupted_verts['right_side'].append(v.id_)

        if max_left_vert is None and max_right_vert is None:
            for v in vertices:
                v_in_side = left_of(pos, pos_next, v.point_)
                if v.convex_ and not v.is_disrupted:
                    dist = point_line_dist(pos, pos_next, v.point_)
                    if v_in_side > 0:
                        if dist > max_left_dist:
                            max_left_dist = dist
                            max_left_vert = v
                    elif v_in_side < 0:
                        if dist > max_right_dist:
                            max_right_dist = dist
                            max_right_vert = v
                elif v.is_disrupted:
                    if v_in_side > 0:
                        disrupted_verts['left_side'].append(v.id_)
                    elif v_in_side < 0:
                        disrupted_verts['right_side'].append(v.id_)

        if len(disrupted_verts['left_side']) > 0 and len(disrupted_verts['right_side']) == 0:       # 往左无解
            max_left_vert, max_left_dist = None, -1
        elif len(disrupted_verts['left_side']) == 0 and len(disrupted_verts['right_side']) > 0:       # 往右无解
            max_right_vert, max_right_dist = None, -1
        return max_left_vert, max_left_dist, max_right_vert, max_right_dist

    def cal_guidance_vertice(self, pos, pos_next, intersect_ob):
        l_vert, l_dist, r_vert, r_dist = self.determin_cand_verts(pos, pos_next, intersect_ob)
        if l_vert is not None and r_vert is not None:
            if l_dist <= r_dist:
                return l_vert, l_dist
            else:
                return r_vert, r_dist
        elif l_vert is not None and r_vert is None:
            return l_vert, l_dist
        elif l_vert is None and r_vert is not None:
            return r_vert, r_dist
        else:
            return None, -1, -1

    def gos(self, pos, pos_next, interseted_obs):
        t1 = time.time()
        if interseted_obs:  # 前方有障碍物阻挡
            max_dist_vertex = None
            max_dist = 0.0
            for intersect_ob in interseted_obs:
                p_vertex, new_dist = self.cal_guidance_vertice(pos, pos_next, intersect_ob)
                if p_vertex is None:
                    t2 = time.time()
                    self.search_time += t2 - t1
                    return None
                if new_dist > max_dist:
                    max_dist = new_dist
                    max_dist_vertex = p_vertex
                if (pos[0], pos[1]) == self.s_start and (self.s_start, self.s_goal) in self.searched:
                    gp = p_vertex.point_
                    p2l_dist = point_line_dist(np.array(self.s_start), np.array(self.s_goal), gp)
                    if p2l_dist > self.searched[(self.s_start, self.s_goal)][2]:
                        self.searched[(self.s_start, self.s_goal)] = (gp[0], gp[1], p2l_dist)

            if max_dist_vertex is None:
                t2 = time.time()
                self.search_time += t2 - t1
                return None
            else:
                gp = max_dist_vertex.point_
                if (pos[0], pos[1]) == self.s_start and (pos_next[0], pos_next[1]) == self.s_goal:
                    self.global_guide_point = [gp[0], gp[1], max_dist]
                    self.searched[(self.s_start, self.s_goal)] = (gp[0], gp[1], max_dist)
                t2 = time.time()
                self.search_time += t2 - t1
                return tuple(gp)
        else:
            t2 = time.time()
            self.search_time += t2 - t1
            return tuple(pos_next)

    def plot_path(self, path, name="GOS"):
        from matplotlib import rcParams
        fig = plt.figure(0)
        fig_size = (10 * 1., 7 * 1.)
        fig.set_size_inches(fig_size[0], fig_size[1])
        ax = fig.add_subplot(1, 1, 1)
        colors = rcParams['axes.prop_cycle'].by_key()['color']  # 获取默认颜色循环
        for i in range(len(self.poly_obs)):
            color = colors[i % len(colors)]
            draw_polygon_2d(ax, self.poly_obs[i].vertices_pos, fc_color=color, ec_color=color, alpha=1)
            draw_polygon_2d(ax, self.inflated_obs[i].vertices_pos, fc_color='none', ec_color=color, alpha=0.9)

        px = [x[0] for x in path]
        py = [x[1] for x in path]
        # spx = [x[0] for x in self.searched]
        # spy = [x[1] for x in self.searched]
        # ax.scatter(spx[:], spy[:], marker='*', color='green')
        plt.plot(px, py, linewidth=2, color='blue', zorder=3)
        plt.plot(self.s_start[0], self.s_start[1], marker='s', color='b', zorder=4)
        plt.plot(self.s_goal[0], self.s_goal[1], marker='s', color='r', zorder=4)
        ax.scatter(px[:], py[:], marker='o', color='red')

        plt.title(name)
        plt.axis("equal")
        plt.show()


def main():
    # envs_type = ["small", "large_maze", "large", "large_corridor", "symmetric", "large_obs", "small_indoor"]
    envs = env.Env(envs_type[3])
    x_range, y_range = (-300, envs.xDim + 300), (0, envs.yDim)
    poly_obs = envs.poly_obs

    # 小范围场景， 50*30
    # radius = 0.5
    # inflation = 0.1
    # agent_starts = [[47, 3], [47, 4], [47, 5], [47, 6], [47, 7], [47, 8], [47, 9], [47, 10], [47, 11], [47, 12],
    #                 [47, 13], [47, 14], [47, 15], [47, 16], [47, 17], [47, 18], [47, 19], [47, 20], [47, 21], [47, 22],
    #                 [47, 23], [47, 24], [47, 25], [47, 26], [47, 27]]
    # agents_goal = [[3, 3], [3, 4], [3, 5], [3, 6], [3, 7], [3, 8], [3, 9], [3, 10], [3, 11], [3, 12], [3, 13], [3, 14],
    #                [3, 15], [3, 16], [3, 17], [3, 18], [3, 19], [3, 20], [3, 21], [3, 22], [3, 23], [3, 24], [3, 25],
    #                [3, 26], [3, 27]]

    # 大范围的走廊场景，读取json文件栅格绘制, 800m*200m
    radius = 5.0
    inflation = min(radius / 2, 1.2)
    agent_starts = [[785.0, 12.0], [785.0, 27.0], [785.0, 42.0], [785.0, 57.0], [785.0, 72.0], [785.0, 87.0], [50, 60],
                    [785.0, 102.0], [785.0, 117.0], [785.0, 132.0], [785.0, 147.0], [785.0, 162.0], [785.0, 177.0]]
    agents_goal = [[15, 12.0], [15, 27.0], [47, 42.0], [47, 57.0], [47, 72.0], [15, 87.0], [700, 150], [15, 102.0],
                   [15, 117.0], [15, 132.0], [70, 147.0], [70, 162.0], [70, 177.0]]

    # 千米级大范围
    # radius = 2.5
    # inflation = 1.
    # pos_sample = [[6100.0, 3700.0], [6100.0, 1125.0], [6100.0, 3925.0], [6100.0, 4900.0], [6100.0, 200.0], [6100.0, 4225.0], [6100.0, 4550.0], [6100.0, 1950.0], [6100.0, 4725.0], [6100.0, 1100.0], [6100.0, 400.0], [6100.0, 500.0], [6100.0, 2900.0], [6100.0, 2175.0], [6100.0, 4875.0], [6100.0, 4275.0], [6100.0, 225.0], [6100.0, 4300.0], [6100.0, 3050.0], [6100.0, 4925.0], [6100.0, 3625.0], [6100.0, 450.0], [6100.0, 2850.0], [6100.0, 3675.0], [6100.0, 2275.0], [6100.0, 1425.0], [6100.0, 750.0], [6100.0, 3375.0], [6100.0, 4025.0], [6100.0, 550.0], [6100.0, 2650.0], [6100.0, 1675.0], [6100.0, 3300.0], [6100.0, 4200.0], [6100.0, 3875.0], [6100.0, 2200.0], [6100.0, 3825.0], [6100.0, 2250.0], [6100.0, 3900.0], [6100.0, 1225.0], [6100.0, 4125.0], [6100.0, 4850.0], [6100.0, 4525.0], [6100.0, 4425.0], [6100.0, 700.0], [6100.0, 3075.0], [6100.0, 3250.0], [6100.0, 900.0], [6100.0, 775.0], [6100.0, 150.0], [6100.0, 3275.0], [6100.0, 4800.0], [6100.0, 975.0], [6100.0, 3175.0], [6100.0, 3800.0], [6100.0, 4575.0], [6100.0, 1075.0], [6100.0, 2350.0], [6100.0, 4975.0], [6100.0, 2950.0], [6100.0, 1700.0], [6100.0, 1975.0], [6100.0, 425.0], [6100.0, 2225.0], [6100.0, 350.0], [6100.0, 2300.0], [6100.0, 2550.0], [6100.0, 3550.0], [6100.0, 1150.0], [6100.0, 1925.0], [6100.0, 475.0], [6100.0, 725.0], [6100.0, 4450.0], [6100.0, 925.0], [6100.0, 2000.0], [6100.0, 2875.0], [6100.0, 3450.0], [6100.0, 1400.0], [6100.0, 3850.0], [6100.0, 1525.0], [6100.0, 1175.0], [6100.0, 4000.0], [6100.0, 1025.0], [6100.0, 125.0], [6100.0, 1250.0], [6100.0, 2450.0], [6100.0, 825.0], [6100.0, 600.0], [6100.0, 2150.0], [6100.0, 625.0], [6100.0, 1650.0], [6100.0, 2575.0], [6100.0, 1350.0], [6100.0, 4775.0], [6100.0, 1825.0], [6100.0, 3600.0], [6100.0, 3425.0], [6100.0, 4475.0], [6100.0, 1500.0], [6100.0, 2750.0]]
    # agent_starts = pos_sample[:100]
    # agents_goal = [[-150, agent_starts[i][1]] for i in range(len(agent_starts))]

    all_length = 0.
    all_time = 0.
    time_data = []
    dist_data = []
    pos_num = len(agent_starts)
    for i in range(0, pos_num):
        idx = i
        s_goal, s_start = (agent_starts[idx][0], agent_starts[idx][1]), (agents_goal[idx][0], agents_goal[idx][1])
        # s_start, s_goal = (2300.66, 4466.02), (4335.71, 4973.88)
        gos = GOS()
        gos.init_rob_info(s_start, s_goal, radius, inflation=inflation)
        gos.init_env(poly_obs, x_range, y_range)

        t1 = time.time()
        path = gos.global_search()
        t2 = time.time()
        all_length += gos.length
        all_time += t2 - t1
        time_data.append(t2 - t1)
        dist_data.append(gos.length)
        print('slove time: ', t2 - t1)
        print('check time: ', gos.check_time)
        print('length: ', gos.length)
        print(len(path), path)
        info = "(Cost: " + str(round(t2 - t1, 5)) + " sec; " + str(round(gos.length, 2)) + " m)"
        envs.plot_env(gos, "GOS" + info)
    print('all slove time: ', all_time)
    print('all length: ', all_length)
    print(time_data)
    print(dist_data)


if __name__ == '__main__':
    main()
